import pandas as pd
import statsmodels.api as sm
import numpy as np
import matplotlib.pyplot as plt


"""
num_award是需要预测的质量变量
"""

#read data
data = pd.read_csv('./data/StudentData.csv', delimiter=',',header=0)

print ('----- data head -----')
print( data.head())
print( '----- data description -----')
print( data.describe())
histData = []
uniqProgs = sorted(data['prog'].unique()) # [1, 2, 3]
for elem in uniqProgs:
    # 将学生表中prog相同的学生的num_awards存在同一个表格中
    histData.append(data[data['prog'] == elem]['num_awards'].values)

# plotting histogram in order to see
# 密度直方图, 反应每一个prog-y关于num_award-x的分布情况
plt.hist(tuple(histData),bins=10, density=True,histtype='bar',label= map(lambda x: 'Prog '+ str(x),uniqProgs))
plt.legend()
plt.ylabel('Count')
plt.title('Histogram for each program')
plt.show()

# adding dummy variables in order to handle categorical data in prog
prog_dummies = pd.get_dummies(data['prog']).rename(columns=lambda x: 'prog_' + str(x))
dataWithDummies = pd.concat([data, prog_dummies], axis=1)
dataWithDummies .drop(['prog', 'prog_3'], inplace=True, axis=1)
dataWithDummies = dataWithDummies .applymap(np.int)

print (dataWithDummies.head())

# applying poisson regression on data
# assuming variables are independent to each other
feat_cols = ['math', 'prog_1', 'prog_2']
X = [elem for elem in dataWithDummies[feat_cols].values] # (200, 4)
# adding constant to adding bias
X = sm.add_constant(X, prepend=False)
Y = [elem for elem in dataWithDummies['num_awards'].values]

# building the model
trainTestSplitNum = 160 # 总数为200, 这个是训练集和测试集的分割点
print(Y[:trainTestSplitNum], X[:trainTestSplitNum])
poisson_mod = sm.Poisson(Y[:trainTestSplitNum], X[:trainTestSplitNum])
poisson_res = poisson_mod.fit(method="newton")
print(poisson_res.summary())


# testing the model
predVals = poisson_res.predict(X[trainTestSplitNum:])

plt.plot(range(len(Y[trainTestSplitNum:])), Y[trainTestSplitNum:], 'r*-', range(len(Y[trainTestSplitNum:])), predVals, 'bo-')
plt.title('Train dataset Real vs. Predicted Values')
plt.legend(['Real Values', 'Predicted Values'])
plt.show()


print("隔点采样")
poisson_mod = sm.Poisson(Y[::2], X[::2])
poisson_res = poisson_mod.fit(method="newton")
print(poisson_res.summary())


# testing the model
predVals = poisson_res.predict(X[1::2])

plt.plot(range(len(Y[1::2])), Y[1::2], 'r*-', range(len(Y[1::2])), predVals, 'bo-')
plt.title('Train dataset Real vs. Predicted Values')
plt.legend(['Real Values', 'Predicted Values'])
plt.show()